package LibDL.supporting;

import LibDL.core.*;
import net.librec.math.structure.DenseMatrix;
import net.librec.math.structure.DenseVector;
import net.librec.math.structure.SequentialAccessSparseMatrix;
import net.librec.math.structure.SparseTensor;
import net.librec.recommender.item.KeyValue;
//import net.librec.recommender.item.RecommendedItemList;
import net.librec.recommender.item.RecommendedList;

import java.util.Arrays;
import java.util.Iterator;
import java.util.*;


public class DataConverter {

    public static SequentialAccessSparseMatrix conert1d(Tensor tensor){
        return null;
    }

    public static SparseTensor convert(Tensor tensor) {
        if (!tensor.is_sparse()) {
            throw new UnsupportedOperationException("Need a sparse tensor");
        }
        if (!tensor.is_coalesced()) {
            tensor = tensor.coalesce();
        }
        StdVectorDouble values = tensor.values().tolist_double();
        StdVectorLong indices = tensor.indices().tolist_long();
        int ndim = tensor.indices().size(0);
        List<Integer>[] keys = new List[ndim];
        for (int i = 0; i < ndim; i++) {
            keys[i] = new ArrayList<>();
            for (int j = 0; j < values.size(); j++) {
                keys[i].add(indices.get(i * values.size() + j));
            }
        }
        int[] dims = new int[tensor.dim()];
        for (int i = 0; i < dims.length; i++) {
            dims[i] = tensor.size(i);
        }
        return new SparseTensor(dims, keys, values);
    }

    public static Tensor convert(SparseTensor tensor) {
        StdVectorLong indices = new StdVectorLong();
        int ndim = tensor.numDimensions;
        for (int i = 0; i < ndim; i++) {
            indices.addAll(tensor.ndKeys[i]);
        }

        StdVectorDouble values = new StdVectorDouble(tensor.values);
        return functional.sparse_coo_tensor(
                functional.tensor(indices).view(new StdVectorLong(Arrays.asList(ndim, -1))),
                functional.tensor(values), new StdVectorLong(tensor.dimensions)
        );
    }


    public static DenseMatrix tensorToMat(Tensor tensor) {
        assert (tensor.dim() == 2);
        tensor = tensor.to(new Device("cpu"));
        tensor = tensor.to(Dtype.FLOAT64);
        double[][] naiveBuffer = new double[tensor.size(0)][tensor.size(1)];
        StdVectorDouble doubles = tensor.tolist_double();
        Iterator<Double> it = doubles.iterator();
        for (int i = 0; i < tensor.size(0); i++) {
            for (int j = 0; j < tensor.size(1); j++) {
                naiveBuffer[i][j] = it.next();
            }
        }
        return new DenseMatrix(naiveBuffer);
    }

    public static Tensor matToTensor(DenseMatrix mat) {
        StdVectorDouble doubles = new StdVectorDouble();
        for (int i = 0; i < mat.rowSize(); i++) {
            for (int j = 0; j < mat.columnSize(); j++) {
                doubles.add(mat.get(i, j));
            }
        }
        Tensor r = functional.as_tensor(doubles);
        try (NoGrad noGrad = new NoGrad()) {
            r = r.view(new StdVectorLong(Arrays.asList(mat.rowSize(), mat.columnSize())));
        }
        return r;
    }

    /*
     * RecommendedList r = tensorToRl(t);
     * evalContext.setRecommendedList(r);
     * RecommenderEvaluator eval = new NormalizedDCGEvaluator();
     * eval.evaluate(evalContext);
     * */
    public static RecommendedList tensorToRl(Tensor item, long flag) {
        assert (item.dim() == 2);
        RecommendedList ret = new RecommendedList(item.size(0));
        item = item.to(new Device("cpu"));

        StdVectorLong key = item.tolist_long();
        Iterator it = key.iterator();
        for (int i = 0, userNum = item.size(0), len = item.size(1); i < userNum; i++) {
            ArrayList<KeyValue<Integer, Double>> al = new ArrayList<KeyValue<Integer, Double>>(len);
            for (int j = 0; j < len; j++) {
                long itemId = (long) it.next();
                if (itemId != flag) {
                    al.add(new KeyValue<Integer, Double>((int) itemId, 0.0));
                }
            }
            ret.addList(al);
        }
        return ret;
    }

    public static RecommendedList tensorToRl(Tensor item) {
        return tensorToRl(item, -1);
    }

    public static RecommendedList tensorToRl(Tensor item, Tensor rating, long flag) {
        assert (item.dim() == 2 && rating.dim() == 2);
        RecommendedList ret = new RecommendedList(item.size(0));
        item = item.to(new Device("cpu"));
        rating = rating.to(new Device("cpu"));

        StdVectorLong key = item.tolist_long();
        StdVectorDouble val = rating.tolist_double();
        Iterator it = key.iterator();
        Iterator<Double> it_ = val.iterator();
        for (int i = 0, userNum = item.size(0), len = item.size(1); i < userNum; i++) {
            ArrayList<KeyValue<Integer, Double>> al = new ArrayList<KeyValue<Integer, Double>>(len);
            for (int j = 0; j < len; j++) {
                long itemId = (long) it.next();
                Double d = it_.next();
                if (itemId != flag) {
                    al.add(new KeyValue<Integer, Double>((int) itemId, d));
                }

            }
            ret.addList(al);
        }
        return ret;
    }

    public static RecommendedList tensorToRl(Tensor item, Tensor rating) {
        return tensorToRl(item, rating, -1);
    }


    public static KeyValue<Tensor, Tensor> rlToTensorPair(RecommendedList list, long flag) {
        int userNum = list.size();
        int maxLen = 0;
        StdVectorLong longs = new StdVectorLong();
        StdVectorDouble doubles = new StdVectorDouble();
        for (int i = 0; i < userNum; i++) {
            List<KeyValue<Integer, Double>> kvl = list.getKeyValueListByContext(i);
            maxLen = Math.max(kvl.size(), maxLen);
        }
        for (int i = 0; i < userNum; i++) {
            List<KeyValue<Integer, Double>> kvl = list.getKeyValueListByContext(i);
            int j = 0;
            for (; j < kvl.size(); j++) {
                longs.add(Integer.valueOf(kvl.get(j).getKey()).longValue());
                doubles.add(kvl.get(j).getValue());
            }
            for (; j < maxLen; j++) {
                longs.add(flag);
                doubles.add(0.0);
            }
        }
        StdVectorLong shape = new StdVectorLong(new long[]{userNum, maxLen});
        return new KeyValue<Tensor, Tensor>(functional.as_tensor(longs).reshape(shape), functional.as_tensor(doubles).reshape(shape));
    }

    public static KeyValue<Tensor, Tensor> rlToTensorPair(RecommendedList list) {
        return rlToTensorPair(list, -1);
    }

    public static Tensor rlToTensor(RecommendedList list, long flag) {
        int userNum = list.size();
        int maxLen = 0;
        StdVectorLong longs = new StdVectorLong();
        for (int i = 0; i < userNum; i++) {
            List<KeyValue<Integer, Double>> kvl = list.getKeyValueListByContext(i);
            maxLen = Math.max(kvl.size(), maxLen);
        }
        for (int i = 0; i < userNum; i++) {
            List<KeyValue<Integer, Double>> kvl = list.getKeyValueListByContext(i);
            int j = 0;
            for (; j < kvl.size(); j++) {
                longs.add(Integer.valueOf(kvl.get(j).getKey()).longValue());
            }
            for (; j < maxLen; j++) {
                longs.add(flag);
            }
        }
        return functional.as_tensor(longs).reshape(new StdVectorLong(new long[]{userNum, maxLen}));
    }

    public static Tensor rlToTensor(RecommendedList list) {
        return rlToTensor(list, -1);
    }

    public static Tensor vector2Tensor(DenseVector vector) {
        return vector2Tensor(vector, new TensorOptions(Dtype.FLOAT32));
    }
    public static Tensor vector2Tensor(DenseVector vector, TensorOptions options) {
        int n = vector.size();
        StdVectorDouble stdVectorDouble = new StdVectorDouble();
        for (int i = 0; i < n; i++) {
            stdVectorDouble.add(vector.get(i));
        }
        return functional.as_tensor(stdVectorDouble, options);
    }
}
